// Copyright (c) 2014 Morten Bakkedal
// This code is published under the MIT License.

using System;
using System.Collections.Generic;
using System.Diagnostics;

using FuncLib.DualFunctions.Compilation;
using FuncLib.Functions;

namespace FuncLib.DualFunctions
{
	/// <summary>
	/// Experimental DualNumber like Function. This class is highly experimental.
	/// </summary>
	[Serializable]
	public abstract partial class DualFunction : Function
	{
		private static int nextFunctionIndex;
		private DualEvaluator evaluator; // FIXME remove
		private CompiledDualFunction compiledFunction;

		static DualFunction()
		{
			nextFunctionIndex = 0;
		}

		public DualFunction()
		{
			FunctionIndex = nextFunctionIndex++;
		}

		protected override double ComputeValue(IEvaluator evaluator)
		{
			//try
			//{
			string identifier;
			return CompiledValue(evaluator, out identifier)[identifier];
			//}
			//catch
			//{
			//    return double.NaN;
			//}

			//return Compile().Value(evaluation);
		}

		protected override Function ComputeDerivative(Variable variable)
		{
			if (variable == null)
			{
				throw new ArgumentNullException("variable");
			}

			if (!(variable is DualInnerVariable))
			{
				// A DualFunction can only depend on a DualVariable (or rather its inner DualInnerVariable instance).
				return 0.0;
			}

			return ComputeDerivative(((DualInnerVariable)variable).OuterVariable);

			//return Compile().Derivative(variable);
		}

		private Function ComputeDerivative(DualVariable variable)
		{
			if (!DualVariable.CheckFirstOrder(variable))
			{
				throw new DualVariableOrderExceededException(variable);
			}

			//if (!gradient.Contains(variable)) // FIXME skip this check? won't work, can't afford to keep all gradient/hessian info in memory for the entire recursive structure of functions
			//{
			//    return 0.0;
			//}

			// Return a wrapper when one of the derivatives of this DualFunction instances is used like an ordinary Function.
			return new DualFunctionFirstDerivative(this, variable);
		}

		private Dictionary<string, double> CompiledValue(IEvaluator evaluation, out string identifier)
		{
			if (evaluator == null)
			{
				evaluator = new DualCodeGenerator(this).Compile();
			}

			identifier = evaluator.Identifier;

			Dictionary<int, double> x = new Dictionary<int, double>();
			Dictionary<string, double> y = new Dictionary<string, double>();

			foreach (VariableAssignment assignment in evaluation.Point)
			{
				DualInnerVariable variable = assignment.Variable as DualInnerVariable;
				if (variable != null)
				{
					// Only DualVariable instances are passed to the invoked method below.
					x[variable.OuterVariable.FunctionIndex] = assignment.Value;
				}
			}

			evaluator.Value(x, y);

			return y;
		}
		
		/// <summary>
		/// Compiles this function for fast evaluation. This is also done implicitly if any of the methods <see cref="Value" />
		/// or <see cref="Derivative" /> are called. The resulting compiled function is cached if called repeatably.
		/// Use <see cref="DualCompiler" /> to compile a group of functions.
		/// </summary>
		public CompiledDualFunction Compile()
		{
			if (compiledFunction != null)
			{
				compiledFunction = DualCompiler.Compile(this);
			}

			return compiledFunction;
		}

		internal abstract GeneratedFunction Compile(DualCodeGenerator generator);

		internal abstract IEnumerable<DualFunction> GetDependencies();

		public static bool TryCast(Function function, out DualFunction dualFunction)
		{
			if (function is DualFunction)
			{
				dualFunction = (DualFunction)function;
				return true;
			}

			if (function is DualInnerVariable)
			{
				dualFunction = ((DualInnerVariable)function).OuterVariable;
				return true;
			}

			dualFunction = null;
			return false;
		}

		/// <summary>
		/// Disables DualFunction arithmetics. Ordinary Function arithmetics will be performed when the returned object is used.
		/// Calling Disable on a DualVariable completely disable DualFunction arithmetics. Useful to avoid DualFunction overhead
		/// when small computations 
		/// </summary>
		public static Function Disable(Function function)
		{
			if (function is DualVariable)
			{
				// Since Function arithmetics is requested, the DualInnerVariable object will make more sence here.
				return Disable(((DualVariable)function).InnerVariable);
			}

			// Can't perform DualFunction arithmetics when wrapped in this object.
			return new DualDisableFunction(function);
		}

		/// <summary>
		/// A unique auto-incremented number assigned to each <see cref="DualFunction" /> instance. Used when referencing functions
		/// in the generated code, and to avoid computing the full symmetric Hessian, but only the lower part of it.
		/// </summary>
		public int FunctionIndex // FIXME internal?
		{
			get;
			private set;
		}

		public Function TestFunction // FIXME Remove this completely when fully tested!!
		{
			get;
			protected set;
		}

		[Serializable]
		private class DualFunctionFirstDerivative : Function
		{
			private DualFunction function;
			private DualVariable variable1;

			public DualFunctionFirstDerivative(DualFunction function, DualVariable variable1)
			{
				this.function = function;
				this.variable1 = variable1;
			}

			protected override double ComputeValue(IEvaluator evaluator)
			{
				string identifier;
				double value;
				if (!function.CompiledValue(evaluator, out identifier).TryGetValue(identifier + "_" + variable1.FunctionIndex, out value))
				{
					// FIXME Assume no dependency on this variable. Make better test of this.
					value = 0.0;
				}

				return value;
			}
			
			protected override Function ComputeDerivative(Variable variable2)
			{
				if (variable2 == null)
				{
					throw new ArgumentNullException("variable2");
				}

				if (!(variable2 is DualInnerVariable))
				{
					// A DualFunction can only depend on a DualVariable (or rather its inner DualInnerVariable instance).
					return 0.0;
				}

				return ComputeDerivative(((DualInnerVariable)variable2).OuterVariable);
			}

			private Function ComputeDerivative(DualVariable variable2)
			{
				if (!DualVariable.CheckSecondOrder(variable1, variable2))
				{
					throw new DualVariableOrderExceededException(variable1, variable2);
				}

				//if (!function.hessian.Contains(variable1, variable2)) // FIXME skip this check? won't work, can't afford to keep all gradient/hessian info in memory for the entire recursive structure of functions
				//{
				//    return 0.0;
				//}

				if (variable1.FunctionIndex > variable2.FunctionIndex)
				{
					// Reverse order of differentiation.
					return new DualFunctionSecondDerivative(function, variable2, variable1);
				}

				return new DualFunctionSecondDerivative(function, variable1, variable2);
			}
		}

		[Serializable]
		private class DualFunctionSecondDerivative : Function
		{
			private DualFunction function;
			private DualVariable variable1, variable2;

			public DualFunctionSecondDerivative(DualFunction function, DualVariable variable1, DualVariable variable2)
			{
				this.function = function;
				this.variable1 = variable1;
				this.variable2 = variable2;
			}

			protected override double ComputeValue(IEvaluator evaluator)
			{
				string identifier;
				double value;
				if (!function.CompiledValue(evaluator, out identifier).TryGetValue(identifier + "_" + variable1.FunctionIndex + "_" + variable2.FunctionIndex, out value))
				{
					// FIXME Assume no dependency on this variable. Make better test of this.
					value = 0.0;
				}

				return value;
			}
			
			protected override Function ComputeDerivative(Variable variable3)
			{
				if (variable3 == null)
				{
					throw new ArgumentNullException("variable3");
				}

				if (!(variable3 is DualInnerVariable))
				{
					// A DualFunction can only depend on a DualVariable (or rather its inner DualInnerVariable instance).
					return 0.0;
				}

				throw new DualVariableOrderExceededException(variable1, variable2, ((DualInnerVariable)variable3).OuterVariable);
			}
		}

		[Serializable]
		private class DualDisableFunction : Function
		{
			public Function f;

			public DualDisableFunction(Function f)
			{
				this.f = f;
			}

			protected override double ComputeValue(IEvaluator evaluator)
			{
				return f.Value(evaluator);
			}

			protected override Function ComputeDerivative(Variable variable)
			{
				return f.Derivative(variable);
			}
		}
	}
}
